# face_enroll.py - 人脸录入和训练模块
import cv2
import os
import numpy as np
from PIL import Image
from config import add_user_name, get_user_names
from face_recognizer import CASCADE_PATH, DATA_PATH, MODEL_PATH

def enroll_face(user_name):
    """录入人脸样本"""
    # 获取或创建用户ID
    user_id = add_user_name(user_name)
    
    # 创建用户专属目录
    user_dir = os.path.join(DATA_PATH, user_name)
    os.makedirs(user_dir, exist_ok=True)
    
    # 调用摄像头
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        raise Exception("无法打开摄像头")

    # 加载人脸检测分类器
    face_detector = cv2.CascadeClassifier(CASCADE_PATH)

    # 初始化样本计数器
    count = 0
    max_count = 249  # 最多采集100张样本

    print(f"\n开始为 {user_name}(ID:{user_id}) 采集面部样本...")
    print("请面向摄像头，保持不同表情和角度")
    print("按下ESC键退出采集")

    while count < max_count:
        # 读取摄像头帧
        success, img = cap.read()
        if not success:
            break

        # 转换为灰度图像
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 检测人脸
        faces = face_detector.detectMultiScale(gray, 1.3, 5)

        # 遍历检测到的人脸
        for (x, y, w, h) in faces:
            # 在彩色图像上绘制蓝色矩形框
            cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)

            # 样本数+1
            count += 1

            # 保存人脸ROI（灰度图）到用户目录
            cv2.imwrite(
                os.path.join(user_dir, f"{user_name}.{user_id}.{count}.jpg"),
                gray[y:y + h, x:x + w]
            )
            
            # 显示采集进度
            cv2.putText(img, f"Collected: {count}/{max_count}", (10, 30), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

        # 显示实时画面
        cv2.imshow('Face Data Collection', img)
        
        # 按ESC退出
        if cv2.waitKey(1) == 27:
            break

    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    
    print(f"\n成功为 {user_name} 采集了 {count} 张面部样本")
    return count

def train_model():
    """训练人脸识别模型"""
    # 初始化人脸检测器和识别器
    detector = cv2.CascadeClassifier(CASCADE_PATH)
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    
    face_samples = []
    ids = []
    
    # 获取所有用户
    users = get_user_names()
    
    for user_name in users:
        user_id = add_user_name(user_name)  # 确保用户ID正确
        user_dir = os.path.join(DATA_PATH, user_name)
        
        if not os.path.exists(user_dir):
            continue
            
        # 获取用户的所有图像文件
        image_files = [f for f in os.listdir(user_dir) if f.lower().endswith('.jpg')]
        
        for image_file in image_files:
            image_path = os.path.join(user_dir, image_file)
            
            # 读取图像并转换为灰度
            img = Image.open(image_path).convert('L')
            img_np = np.array(img, 'uint8')
            
            # 人脸检测
            faces = detector.detectMultiScale(img_np)
            for (x, y, w, h) in faces:
                face_samples.append(img_np[y:y + h, x:x + w])
                ids.append(user_id)
    
    if len(face_samples) == 0:
        return False, 0

    # 训练模型
    recognizer.train(face_samples, np.array(ids))
    
    # 保存模型
    os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
    recognizer.save(MODEL_PATH)
    
    return True, len(face_samples)